import os
import random
import time

from dataset import load_dataset
from lazyselect import lazyselect_once, lazyselect
from utils import brief, save_variable

if __name__ == '__main__':
    run_num = 1000
    name = 'lazyselect_analysis_with_fail'

    dataset_size = 50000
    r_exps = [i / 8 for i in range(1, 9)]
    rand = random.Random()
    rand.seed(5)

    print('r_exp:', r_exps)

    s = load_dataset(os.path.join('datasets', f'dataset-{str(dataset_size)}.txt'))
    print('dataset size:', len(s))

    k_list = rand.choices(range(1, dataset_size + 1), k=run_num)
    # print('k_list:', brief(k_list), ', len =', len(k_list))
    avg_consumes = []
    success_rates = []
    fail_rates = {
        'k-th element is not in P': [],
        '|p| >= 4 * n ^ sampling_ratio': []
    }
    for r_exp in r_exps:
        print('sampling_ratio:', r_exp)
        success_count = 0
        why_fail_count = {
            'k-th element is not in P': 0,
            '|p| >= 4 * n ^ sampling_ratio': 0
        }
        start = time.time()
        for k in k_list:
            result = lazyselect_once(s, k, r_exp)
            # if result is not None:
            #     success_count += 1
            if isinstance(result, str):
                why_fail_count[result] += 1
            else:
                success_count += 1
        elapsed = time.time() - start
        elapsed = elapsed / run_num
        avg_consumes.append(elapsed)
        success_rates.append(success_count / run_num)
        fail_rates['k-th element is not in P'].append(
            why_fail_count['k-th element is not in P'] / run_num)
        fail_rates['|p| >= 4 * n ^ sampling_ratio'].append(
            why_fail_count['|p| >= 4 * n ^ sampling_ratio'] / run_num)
        print("Average time consume:", elapsed * 1000, 'ms')
        print('Success count:', success_count)
        print(why_fail_count)
        print('=' * 50)

    output = {
        'sampling_ratios': r_exps,
        'avg_time_consumes': avg_consumes,
        'success_rates': success_rates,
        'fail_rates': fail_rates,
        'run_num': run_num
    }
    os.makedirs('results', exist_ok=True)
    save_variable(output, os.path.join('results', name + '.pkl'))

    print("Finished.")
